package com.canyou.datascope;

import com.baomidou.mybatisplus.core.parser.AbstractJsqlParser;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.experimental.Accessors;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 数据权限 SQL 解析器（ org_id 行级 ）
 *
 * @author fy
 * @date 20/04/17 18:14
 */
@Data
@Accessors(chain = true)
@EqualsAndHashCode(callSuper = true)
public class DataScopeSqlParser extends AbstractJsqlParser {

    private String scopeName;
    private List<Long> scopeIds;

    public DataScopeSqlParser(String scopeName, List<Long> scopeIds) {
        this.scopeName = scopeName;
        this.scopeIds = scopeIds;
    }

    @Override
    public void processSelectBody(SelectBody selectBody) {
        if (selectBody instanceof PlainSelect) {
            processPlainSelect((PlainSelect) selectBody);
        } else if (selectBody instanceof WithItem) {
            WithItem withItem = (WithItem) selectBody;
            if (withItem.getSelectBody() != null) {
                processSelectBody(withItem.getSelectBody());
            }
        } else {
            SetOperationList operationList = (SetOperationList) selectBody;
            if (operationList.getSelects() != null && operationList.getSelects().size() > 0) {
                operationList.getSelects().forEach(this::processSelectBody);
            }
        }
    }

    /**
     * 处理 PlainSelect
     */
    protected void processPlainSelect(PlainSelect plainSelect) {
        processPlainSelect(plainSelect, false);
    }

    /**
     * 处理 PlainSelect
     *
     * @param plainSelect ignore
     * @param addColumn   是否添加数据权限列, select语句中需要
     */
    protected void processPlainSelect(PlainSelect plainSelect, boolean addColumn) {
        FromItem fromItem = plainSelect.getFromItem();
        if (fromItem instanceof Table) {
            Table fromTable = (Table) fromItem;
            plainSelect.setWhere(builderExpression(plainSelect.getWhere(), fromTable));
            if (addColumn) {
                plainSelect.getSelectItems().add(new SelectExpressionItem(new Column(this.scopeName)));
            }
        }
    }

    /**
     * 处理条件
     */
    protected Expression builderExpression(Expression expression, Table table) {
        // 生成字段名
        Expression newExp = getOrgIdExpression(table);
        // 加入判断防止条件为空时生成 "and null" 导致查询结果为空
        if (expression == null) {
            return newExp;
        } else {
            if (expression instanceof OrExpression) {
                return new AndExpression(newExp, new Parenthesis(expression));
            } else {
                return new AndExpression(newExp, expression);
            }
        }
    }

    /**
     * 数据权限字段别名设置
     * <p>tableName.org_id 或 tableAlias.org_id</p>
     *
     * @param table 表对象
     * @return 字段
     */
    protected Column getAliasColumn(Table table) {
        StringBuilder column = new StringBuilder();
        // fix 应对当数据权限字段不在主表上, 传入的字段别名已经包含table alias时, 不需要系统去获取表别名
        if (this.scopeName != null && this.scopeName.contains(StringPool.DOT)) {
            column.append(this.scopeName);
            return new Column(column.toString());
        }
        // 正常获取表别名, 然后拼接上, 如果没有表别名则使用表名
        if (null == table.getAlias()) {
            column.append(table.getName());
        } else {
            column.append(table.getAlias().getName());
        }
        column.append(StringPool.DOT);
        column.append(this.scopeName);
        return new Column(column.toString());
    }

    /**
     * 构建where条件
     *
     * @param table
     * @return
     */
    private Expression getOrgIdExpression(Table table) {
        if (this.scopeIds.size() <= 1) {
            EqualsTo equalsTo = new EqualsTo();
            equalsTo.setLeftExpression(this.getAliasColumn(table));
            equalsTo.setRightExpression(new LongValue(this.scopeIds.get(0)));
            return equalsTo;
        } else {
            final InExpression inExpression = new InExpression();
            inExpression.setLeftExpression(this.getAliasColumn(table));
            final ExpressionList itemsList = new ExpressionList();
            final List<Expression> inValues = new ArrayList<>(scopeIds.size());
            inValues.addAll(scopeIds.stream().map(LongValue::new).collect(Collectors.toSet()));
            itemsList.setExpressions(inValues);
            inExpression.setRightItemsList(itemsList);
            return inExpression;
        }
    }

    @Override
    public void processInsert(Insert insert) {
        throw new UnsupportedOperationException("insert not supported");
    }

    @Override
    public void processDelete(Delete delete) {
        throw new UnsupportedOperationException("delete not supported");
    }

    @Override
    public void processUpdate(Update update) {
        throw new UnsupportedOperationException("update not supported");
    }
}
